{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e785571c",
   "metadata": {},
   "source": [
    "## Markov Chains"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18476b1f",
   "metadata": {},
   "source": [
    "author: Jacob Schreiber <br>\n",
    "contact: jmschreiber91@gmail.com\n",
    "\n",
    "Markov chains are the simplest probabilistic model describing a sequence of observations. Essentially, for an n-th order Markov chain, each observation is modeled as $P(X_{t} | X_{t-1}, ..., X_{t-n})$ and the probability of the entire sequence is the product of these probabilities for each observation. Naturally, the first observation in the sequence cannot be modeled as this conditional distribution and so is usually modeled as a marginal distribution. The remaining $n-1$ observations also cannot be modeled as this full conditional distribution and so are modeled by smaller distributions.\n",
    "\n",
    "These chains are easy to implement and use in pomegranate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e54bf2cd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Populating the interactive namespace from numpy and matplotlib\n",
      "torch      : 1.13.0\n",
      "pomegranate: 1.0.0\n",
      "\n",
      "Compiler    : GCC 11.2.0\n",
      "OS          : Linux\n",
      "Release     : 4.15.0-208-generic\n",
      "Machine     : x86_64\n",
      "Processor   : x86_64\n",
      "CPU cores   : 8\n",
      "Architecture: 64bit\n",
      "\n"
     ]
    }
   ],
   "source": [
    "%pylab inline\n",
    "import seaborn; seaborn.set_style('whitegrid')\n",
    "import torch\n",
    "\n",
    "%load_ext watermark\n",
    "%watermark -m -n -p torch,pomegranate"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac9f587c",
   "metadata": {},
   "source": [
    "### Initialization and Fitting\n",
    "\n",
    "Initializing a Markov chain is simple. If you have fit distributions, you can pass them in and then use the model for inferene. If you do not have fit distributions, you can specify the `k` parameter, which is the number of previous observations that the probability of each observation is conditioned on. pomegranate will automatically construct the first `k-1` distributions as well as the main conditional distribution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "45a08fea",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from pomegranate.markov_chain import MarkovChain\n",
    "\n",
    "model = MarkovChain(k=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0ed3346",
   "metadata": {},
   "source": [
    "The model can then be fit to data using the `fit` function. However, this data must be three dimensional, with the dimensions being `(n_samples, length, dimensions)`. Most other models in pomegranate only use the first two. This does mean that Markov chains can be multivariate but a multivariate model will assume each of the dimensions are independent of each other."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "6421a2a7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "MarkovChain()"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X = torch.tensor([[[1], [0], [0], [1]], \n",
    "                  [[0], [1], [0], [0]],\n",
    "                  [[0], [0], [0], [0]],\n",
    "                  [[0], [0], [0], [1]],\n",
    "                  [[0], [1], [1], [0]]])\n",
    "\n",
    "model.fit(X)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b22e05cd",
   "metadata": {},
   "source": [
    "We can then inspect the distribution and compare them to the data to ensure that they're right."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "3f44c198",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0.8000, 0.2000])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.distributions[0].probs[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "0e9bfd76",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Parameter containing:\n",
       "tensor([[0.5000, 0.5000],\n",
       "        [1.0000, 0.0000]])"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.distributions[1].probs[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "cc6e0a3e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Parameter containing:\n",
       "tensor([[[1.0000, 0.0000],\n",
       "         [0.5000, 0.5000]],\n",
       "\n",
       "        [[1.0000, 0.0000],\n",
       "         [0.5000, 0.5000]]])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.distributions[2].probs[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d21b4c03",
   "metadata": {},
   "source": [
    "If we wanted to fit a multivariate model all we would need to do is increase the size of the second dimension."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a403c81b",
   "metadata": {},
   "source": [
    "### Probability and Log Probability"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1bb609d3",
   "metadata": {},
   "source": [
    "The probability of a sequence under a Markov chain is the product of the probabilities of each observation given the previous observations. Another way of putting this is that the joint probability of a sequence with n observations $P(X_{1} ... X_{n})$ is factorized along a chain and equal to $P(X_{1}) P(X_{2} | X_{1}) \\prod\\limits_{t=3}^{n} P(X_{t} | X_{t-1}, X_{t-2})$ for a third order Markov chain.\n",
    "\n",
    "If you data is multivariate, the probability of each dimension is independent and multiplied together at the end. If you would like dependencies between your dimensions, you should consider encoding a single dimension to include all combinations of observations across the features.\n",
    "\n",
    "We can calculate the probability and log probability just as easily as other models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "aee0caa2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0.2000, 0.2000, 0.2000, 0.2000, 0.2000])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.probability(X)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "b9128d9b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([-1.6094, -1.6094, -1.6094, -1.6094, -1.6094])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.log_probability(X)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0dad7b7f",
   "metadata": {},
   "source": [
    "### Summarization"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "332dcf36",
   "metadata": {},
   "source": [
    "Markov chains can perform summarization of data just like other models but that data has to have the three dimensions mentioned before. Further, each chunk of data that is summarized must have the same length. This means that if you have data with different lengths, you must either summarize them one at a time or bucket the sequences such that each bucket has all the sequences of the same length."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "02a61b78",
   "metadata": {},
   "outputs": [],
   "source": [
    "X = torch.randint(2, size=(30, 8, 1))\n",
    "\n",
    "model.summarize(X[:10])\n",
    "model.summarize(X[10:])\n",
    "model.from_summaries()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6965782",
   "metadata": {},
   "source": [
    "### Sampling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "61c3b5e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "X_sample = model.sample(100000).type(torch.float32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "735921b8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Parameter containing:\n",
       "tensor([[0.4667, 0.5333]])"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.distributions[0].probs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "0f604e1d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0.5332)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X_sample[:, 0, 0].mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "2f32b318",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Parameter containing:\n",
       "tensor([[0.3571, 0.6429],\n",
       "        [0.5000, 0.5000]])"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.distributions[1].probs[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "2d399fed",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0.5039)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X_sample[X_sample[:, 0, 0] == 1, 1, 0].mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "b623d9cf",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Parameter containing:\n",
       "tensor([[[0.4000, 0.6000],\n",
       "         [0.2222, 0.7778]],\n",
       "\n",
       "        [[0.1250, 0.8750],\n",
       "         [0.7500, 0.2500]]])"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.distributions[2].probs[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "6477501d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0.2550)"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X_sample[(X_sample[:, 0, 0] == 1) & (X_sample[:, 1, 0] == 1), 2, 0].mean()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
